package com.itheima.java.ai.langchin4j.config;

import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.DocumentParser;
import dev.langchain4j.data.document.parser.TextDocumentParser;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.memory.chat.ChatMemoryProvider;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.store.embedding.EmbeddingStoreIngestor;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;
import dev.langchain4j.store.memory.chat.ChatMemoryStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.core.io.ClassPathResource;
import org.springframework.core.io.Resource;
import org.springframework.core.io.ResourceLoader;

import java.io.IOException;
import java.io.InputStream;
import java.util.Arrays;
import java.util.List;

@Configuration
public class XiaozhiAgentConfig {
    @Autowired
    private ChatMemoryStore chatMemoryStore;
    
    @Autowired
    private ResourceLoader resourceLoader;

    @Bean

    public   ChatMemoryProvider chatMemoryProviderXiaozhi() {
        return memoryId -> MessageWindowChatMemory.builder()
                .id(memoryId)
                .maxMessages(20)
                .chatMemoryStore(chatMemoryStore)
                .build();

    }
    @Bean
    ContentRetriever contentRetrieverXiaozhi() {
        try {
            //使用Spring ResourceLoader读取classpath下的知识库文档
            DocumentParser parser = new TextDocumentParser();
            
            Resource resource1 = resourceLoader.getResource("classpath:knowledge/医院信息.md");
            Document document1 = loadDocumentFromResource(resource1, parser);

            Resource resource2 = resourceLoader.getResource("classpath:knowledge/科室信息.md");
            Document document2 = loadDocumentFromResource(resource2, parser);

            Resource resource3 = resourceLoader.getResource("classpath:knowledge/神经内科.md");
            Document document3 = loadDocumentFromResource(resource3, parser);

            List<Document> documents = Arrays.asList(document1, document2, document3);

            //使用内存向量存储
            InMemoryEmbeddingStore<TextSegment> embeddingStore = new InMemoryEmbeddingStore<>();
            //使用默认的文档分割器
            EmbeddingStoreIngestor.ingest(documents, embeddingStore);
            //从嵌入存储（EmbeddingStore）里检索和查询内容相关的信息
            return EmbeddingStoreContentRetriever.from(embeddingStore);
        } catch (IOException e) {
            throw new RuntimeException("Failed to load knowledge documents", e);
        }
    }
    
    private Document loadDocumentFromResource(Resource resource, DocumentParser parser) throws IOException {
        try (InputStream inputStream = resource.getInputStream()) {
            return parser.parse(inputStream);
        }
    }
}
